#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
动作规划模块

负责根据意图生成具体的动作计划，包括动作序列和参数。
"""

import os
import openai
from typing import Dict, Any, Optional, List
import json


class ActionPlanner:
    """动作规划器类"""

    def __init__(self, api_key: Optional[str] = None):
        """
        初始化动作规划器

        Args:
            api_key: OpenAI API密钥，如果为None则从环境变量获取
        """
        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
        if not self.api_key:
            raise ValueError("未提供API密钥且环境变量OPENAI_API_KEY未设置")

        # 初始化OpenAI
        openai.api_key = self.api_key
        openai.api_base = "https://dashscope.aliyuncs.com/compatible-mode/v1"

        # 动作计划模板
        self.action_plan_template = {
            "sequence": [],  # 动作序列
            "parameters": {},  # 全局参数
            "constraints": [],  # 约束条件
            "priority": 0,  # 优先级
            "estimated_duration": 0.0,  # 预计持续时间
        }

    def plan(self, intent: Dict[str, Any]) -> Dict[str, Any]:
        """
        根据意图生成动作计划

        Args:
            intent: 意图字典

        Returns:
            Dict[str, Any]: 动作计划
        """
        try:
            # 构建提示词
            prompt = self._build_prompt(intent)

            # 调用API获取响应
            response = openai.ChatCompletion.create(
                model="qwen-omni-turbo",
                messages=[
                    {
                        "role": "system",
                        "content": "你是一个六足机器人的动作规划助手。请根据用户的意图生成具体的动作计划。",
                    },
                    {"role": "user", "content": prompt},
                ],
                temperature=0.7,
                max_tokens=1000,
            )

            # 解析响应
            action_plan = self._parse_response(response.choices[0].message.content)

            return action_plan

        except Exception as e:
            print(f"动作规划出错: {e}")
            return {
                "sequence": [],
                "parameters": {"error": str(e)},
                "constraints": [],
                "priority": 0,
                "estimated_duration": 0.0,
            }

    def _build_prompt(self, intent: Dict[str, Any]) -> str:
        """
        构建提示词

        Args:
            intent: 意图字典

        Returns:
            str: 完整的提示词
        """
        return f"""
请根据以下意图生成动作计划：
意图：{json.dumps(intent, ensure_ascii=False, indent=2)}

请以JSON格式返回以下信息：
1. sequence: 动作序列，每个动作包含：
   - name: 动作名称
   - parameters: 动作参数
   - duration: 预计持续时间
2. parameters: 全局参数
3. constraints: 约束条件列表
4. priority: 优先级（0-10）
5. estimated_duration: 总预计持续时间

示例返回格式：
{{
    "sequence": [
        {{
            "name": "prepare_legs",
            "parameters": {{
                "height": 0.2,
                "stance": "tripod"
            }},
            "duration": 1.0
        }},
        {{
            "name": "move_forward",
            "parameters": {{
                "speed": 0.5,
                "distance": 1.0
            }},
            "duration": 2.0
        }}
    ],
    "parameters": {{
        "max_speed": 0.8,
        "min_height": 0.1
    }},
    "constraints": [
        "maintain_balance",
        "avoid_obstacles"
    ],
    "priority": 5,
    "estimated_duration": 3.0
}}
"""

    def _parse_response(self, response: str) -> Dict[str, Any]:
        """
        解析API响应

        Args:
            response: API返回的响应字符串

        Returns:
            Dict[str, Any]: 解析后的动作计划
        """
        try:
            # 尝试解析JSON
            action_plan = json.loads(response)

            # 验证必要字段
            for key in self.action_plan_template:
                if key not in action_plan:
                    action_plan[key] = self.action_plan_template[key]

            return action_plan

        except json.JSONDecodeError:
            # 如果JSON解析失败，返回错误计划
            return {
                "sequence": [],
                "parameters": {"error": "JSON解析失败"},
                "constraints": [],
                "priority": 0,
                "estimated_duration": 0.0,
            }

    def validate_plan(self, action_plan: Dict[str, Any]) -> bool:
        """
        验证动作计划的有效性

        Args:
            action_plan: 动作计划字典

        Returns:
            bool: 计划是否有效
        """
        try:
            # 检查必要字段
            for key in self.action_plan_template:
                if key not in action_plan:
                    return False

            # 检查动作序列
            if not isinstance(action_plan["sequence"], list):
                return False

            for action in action_plan["sequence"]:
                if not all(k in action for k in ["name", "parameters", "duration"]):
                    return False

            # 检查参数
            if not isinstance(action_plan["parameters"], dict):
                return False

            # 检查约束条件
            if not isinstance(action_plan["constraints"], list):
                return False

            # 检查优先级
            if not isinstance(action_plan["priority"], (int, float)):
                return False

            # 检查预计持续时间
            if not isinstance(action_plan["estimated_duration"], (int, float)):
                return False

            return True

        except Exception:
            return False
